from openai import AsyncOpenAI
from openai.types.beta.threads import TextContentBlock

from llm.entity.metadatas import MessageMetadata
from llm.service import OpenAIService


async def output_invoke_stream(input: str, thread_id: str, consumer):
    print(f"开始反馈回话")
    client = AsyncOpenAI()
    assistant_id = "asst_MOE2NoeUmR9MSdHED98JvlkH"
    run = None
    message = await client.beta.threads.messages.create(
        content=input,
        thread_id=thread_id,
        role="user",
        metadata=MessageMetadata(
            type="3",
        )
    )
    async with client.beta.threads.runs.stream(
            assistant_id=assistant_id,
            thread_id=thread_id,
            metadata={
                "type": "response"
            },
            event_handler=OpenAIService.WebSocketOutputHandler(consumer=consumer)
    ) as stream:
        await stream.until_done()
        print(f'{stream.current_run.id}')
        run = stream.current_run
    return run


async def mark_message(thread_id, run):
    client = AsyncOpenAI()
    messages = await client.beta.threads.messages.list(
        thread_id=thread_id,
    )
    latest_assistant_message = None
    for message in messages.data:
        if message.role == 'assistant' and message.run_id == run.id:
            latest_assistant_message = message
            latest_assistant_message = await client.beta.threads.messages.update(
                message_id=message.id,
                metadata=MessageMetadata(
                    type="1"
                ),
                thread_id=thread_id,
            )
            break
    # print(latest_assistant_message)
    # extracted_values = [item.text.value for item in latest_assistant_message.content if
    #                     isinstance(item, TextContentBlock)]
    # return extracted_values


async def output_invoke(input: str, thread_id: str):
    client = AsyncOpenAI()
    assistant_id = "asst_MOE2NoeUmR9MSdHED98JvlkH"
    message = await client.beta.threads.messages.create(
        content=input,
        thread_id=thread_id,
        role="user",
        metadata=MessageMetadata(
            type="3",
        )
    )
    run = await client.beta.threads.runs.create_and_poll(
        assistant_id=assistant_id,
        thread_id=thread_id,
        metadata={
            "type": "response"
        }
    )
    messages = await client.beta.threads.messages.list(
        thread_id=thread_id,
    )
    latest_assistant_message = None
    for message in messages.data:
        if message.role == 'assistant' and message.run_id == run.id:
            latest_assistant_message = message
            latest_assistant_message = await client.beta.threads.messages.update(
                message_id=message.id,
                metadata=MessageMetadata(
                    type="1"
                ),
                thread_id=thread_id,
            )
            break
    print(latest_assistant_message)
    extracted_values = [item.text.value for item in latest_assistant_message.content if
                        isinstance(item, TextContentBlock)]
    return extracted_values
